﻿// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

using System;
using System.Collections.Generic;
using System.Net;
using System.Threading.Tasks;
using Azure.Core.TestFramework;
using NUnit.Framework;

namespace Azure.AI.FormRecognizer.DocumentAnalysis.Tests
{
    /// <summary>
    /// The suite of tests for the document model methods in the <see cref="DocumentModelAdministrationClient"/> class.
    /// </summary>
    /// <remarks>
    /// These tests have a dependency on live Azure services and may incur costs for the associated
    /// Azure subscription.
    /// </remarks>
    public class DocumentModelAdministrationLiveTests : DocumentAnalysisLiveTestBase
    {
        private static readonly DocumentBuildMode[] s_buildDocumentModelTestCases = new[]
        {
            DocumentBuildMode.Template,
            DocumentBuildMode.Neural
        };

        /// <summary>
        /// Initializes a new instance of the <see cref="DocumentModelAdministrationLiveTests"/> class.
        /// </summary>
        /// <param name="isAsync">A flag used by the Azure Core Test Framework to differentiate between tests for asynchronous and synchronous methods.</param>
        public DocumentModelAdministrationLiveTests(bool isAsync, DocumentAnalysisClientOptions.ServiceVersion serviceVersion)
            : base(isAsync, serviceVersion)
        {
        }

        #region Build

        [RecordedTest]
        public async Task BuildDocumentModelCanAuthenticateWithTokenCredential()
        {
            var client = CreateDocumentModelAdministrationClient(useTokenCredential: true);
            var modelId = Recording.GenerateId();
            var trainingFilesUri = new Uri(TestEnvironment.BlobContainerSasUrl);

            BuildDocumentModelOperation operation = null;

            try
            {
                operation = await client.BuildDocumentModelAsync(WaitUntil.Completed, trainingFilesUri, DocumentBuildMode.Template, modelId);
            }
            finally
            {
                if (operation != null && operation.HasValue)
                {
                    await client.DeleteDocumentModelAsync(modelId);
                }
            }

            // Sanity check to make sure we got an actual response back from the service.
            Assert.AreEqual(modelId, operation.Value.ModelId);
        }

        [RecordedTest]
        [TestCaseSource(nameof(s_buildDocumentModelTestCases))]
        public async Task BuildDocumentModel(DocumentBuildMode buildMode)
        {
            if (buildMode == DocumentBuildMode.Neural && Recording.Mode == RecordedTestMode.Live)
            {
                // Test takes too long to finish running, and seems to cause multiple failures in our
                // live test pipeline. For this reason, this test is ignored when running in Live mode.
                Assert.Ignore();
            }

            var client = CreateDocumentModelAdministrationClient();
            var modelId = Recording.GenerateId();
            var startTime = Recording.UtcNow;
            var trainingFilesUri = new Uri(TestEnvironment.BlobContainerSasUrl);
            var prefix = "subfolder";
            var options = new BuildDocumentModelOptions()
            {
                Description = $"This model was generated by a .NET test.",
                Tags = { { "tag1", "value1" }, { "tag2", "value2" } }
            };

            BuildDocumentModelOperation operation = null;

            try
            {
                operation = await client.BuildDocumentModelAsync(WaitUntil.Started, trainingFilesUri, buildMode, modelId, prefix, options);

                await (buildMode == DocumentBuildMode.Neural
                    ? operation.WaitForCompletionAsync(TimeSpan.FromMinutes(1))
                    : operation.WaitForCompletionAsync());
            }
            finally
            {
                if (operation != null && operation.HasValue)
                {
                    await client.DeleteDocumentModelAsync(modelId);
                }
            }

            Assert.IsTrue(operation.HasValue);

            DocumentModelDetails model = operation.Value;

            Assert.AreEqual(modelId, model.ModelId);
            Assert.AreEqual(options.Description, model.Description);
            Assert.AreEqual(ServiceVersionString, model.ServiceVersion);

            // Add a 4-hour tolerance because model could have been cached before this test.
            Assert.Greater(model.CreatedOn, startTime - TimeSpan.FromHours(4));

            if (_serviceVersion >= DocumentAnalysisClientOptions.ServiceVersion.V2023_07_31)
            {
                Assert.Greater(model.ExpiresOn, model.CreatedOn);
            }
            else
            {
                Assert.IsNull(model.ExpiresOn);
            }

            CollectionAssert.AreEquivalent(options.Tags, model.Tags);

            Assert.AreEqual(1, model.DocumentTypes.Count);

            DocumentTypeDetails documentType = model.DocumentTypes[modelId];

            Assert.IsNull(documentType.Description);
            Assert.AreEqual(buildMode, documentType.BuildMode);

            if (buildMode == DocumentBuildMode.Template)
            {
                CollectionAssert.AreEquivalent(documentType.FieldConfidence.Keys, documentType.FieldSchema.Keys);

                foreach (float confidence in documentType.FieldConfidence.Values)
                {
                    Assert.GreaterOrEqual(confidence, 0f);
                    Assert.LessOrEqual(confidence, 1f);
                }
            }
            else
            {
                Assert.IsEmpty(documentType.FieldConfidence);
            }

            foreach (DocumentFieldSchema fieldSchema in documentType.FieldSchema.Values)
            {
                Assert.IsNull(fieldSchema.Description);
                Assert.IsNull(fieldSchema.Example);
                Assert.IsNull(fieldSchema.Items);
                Assert.IsEmpty(fieldSchema.Properties);
            }

            DocumentFieldSchema merchantSchema = documentType.FieldSchema["Merchant"];
            DocumentFieldSchema quantitySchema = documentType.FieldSchema["Quantity"];

            Assert.AreEqual(DocumentFieldType.String, merchantSchema.Type);
            Assert.AreEqual(DocumentFieldType.Double, quantitySchema.Type);
        }

        [RecordedTest]
        public void BuildDocumentModelThrowsWhenTrainingFilesAreMissing()
        {
            var client = CreateDocumentModelAdministrationClient();
            var modelId = Recording.GenerateId();
            var trainingFilesUri = new Uri(TestEnvironment.BlobContainerSasUrl);
            var prefix = "testfolder"; // folder exists but most training files are missing

            RequestFailedException ex = Assert.ThrowsAsync<RequestFailedException>(async () => await client.BuildDocumentModelAsync(WaitUntil.Started, trainingFilesUri, DocumentBuildMode.Template, modelId, prefix));
            Assert.AreEqual("InvalidRequest", ex.ErrorCode);
        }

        #endregion

        #region Copy

        [RecordedTest]
        [TestCase(true)]
        [TestCase(false)]
        public async Task CopyDocumentModelTo(bool useTokenCredential)
        {
            var client = CreateDocumentModelAdministrationClient(useTokenCredential);
            var modelId = Recording.GenerateId();
            var description = $"This model was generated by a .NET test.";
            var tags = new Dictionary<string, string>() { { "tag1", "value1" }, { "tag2", "value2" } };
            var startTime = Recording.UtcNow;

            await using var disposableModel = await BuildDisposableDocumentModelAsync();

            DocumentModelCopyAuthorization copyAuthorization = await client.GetCopyAuthorizationAsync(modelId, description, tags);
            CopyDocumentModelToOperation operation = null;

            try
            {
                operation = await client.CopyDocumentModelToAsync(WaitUntil.Completed, disposableModel.ModelId, copyAuthorization);
            }
            finally
            {
                if (operation != null && operation.HasValue)
                {
                    await client.DeleteDocumentModelAsync(modelId);
                }
            }

            Assert.IsTrue(operation.HasValue);

            DocumentModelDetails sourceModel = disposableModel.Value;
            DocumentModelDetails model = operation.Value;

            Assert.AreEqual(modelId, model.ModelId);
            Assert.AreEqual(description, model.Description);
            Assert.AreEqual(ServiceVersionString, model.ServiceVersion);

            // Add a 4-hour tolerance because model could have been cached before this test.
            Assert.Greater(model.CreatedOn, startTime - TimeSpan.FromHours(4));
            Assert.Greater(model.ExpiresOn, model.CreatedOn);

            CollectionAssert.AreEquivalent(tags, model.Tags);

            // (TODO) Reenable validation once the following service issue has been fixed: https://github.com/Azure/azure-sdk-for-net/issues/37172
            // AssertDocumentTypeDictionariesAreEquivalent(sourceModel.DocumentTypes, model.DocumentTypes);
        }

        #endregion Copy

        #region Compose

        [RecordedTest]
        [TestCase(true)]
        [TestCase(false)]
        public async Task ComposeDocumentModel(bool useTokenCredential)
        {
            var client = CreateDocumentModelAdministrationClient(useTokenCredential);
            var modelId = Recording.GenerateId();
            var description = $"This model was generated by a .NET test.";
            var tags = new Dictionary<string, string>() { { "tag1", "value1" }, { "tag2", "value2" } };
            var startTime = Recording.UtcNow;

            await using var disposableModel0 = await BuildDisposableDocumentModelAsync(ContainerType.Singleforms);
            await using var disposableModel1 = await BuildDisposableDocumentModelAsync(ContainerType.SelectionMarks);

            var componentModelIds = new string[] { disposableModel0.ModelId, disposableModel1.ModelId };
            ComposeDocumentModelOperation operation = null;

            try
            {
                operation = await client.ComposeDocumentModelAsync(WaitUntil.Completed, componentModelIds, modelId, description, tags);
            }
            finally
            {
                if (operation != null && operation.HasValue)
                {
                    await client.DeleteDocumentModelAsync(modelId);
                }
            }

            Assert.IsTrue(operation.HasValue);

            DocumentModelDetails componentModel0 = disposableModel0.Value;
            DocumentModelDetails componentModel1 = disposableModel1.Value;
            DocumentModelDetails model = operation.Value;

            Assert.AreEqual(modelId, model.ModelId);
            Assert.AreEqual(description, model.Description);
            Assert.AreEqual(ServiceVersionString, model.ServiceVersion);

            // Add a 4-hour tolerance because model could have been cached before this test.
            Assert.Greater(model.CreatedOn, startTime - TimeSpan.FromHours(4));

            if (_serviceVersion >= DocumentAnalysisClientOptions.ServiceVersion.V2023_07_31)
            {
                Assert.Greater(model.ExpiresOn, model.CreatedOn);
            }
            else
            {
                Assert.IsNull(model.ExpiresOn);
            }

            CollectionAssert.AreEquivalent(tags, model.Tags);

            Assert.AreEqual(2, model.DocumentTypes.Count);

            DocumentTypeDetails expectedDocumentType0 = componentModel0.DocumentTypes[componentModel0.ModelId];
            DocumentTypeDetails expectedDocumentType1 = componentModel1.DocumentTypes[componentModel1.ModelId];
            DocumentTypeDetails documentType0 = model.DocumentTypes[componentModel0.ModelId];
            DocumentTypeDetails documentType1 = model.DocumentTypes[componentModel1.ModelId];

            DocumentAssert.AreEqual(expectedDocumentType0, documentType0);
            DocumentAssert.AreEqual(expectedDocumentType1, documentType1);
        }

        [RecordedTest]
        public void ComposeDocumentModelThrowsWhenComponentModelDoesNotExist()
        {
            var client = CreateDocumentModelAdministrationClient();
            var fakeComponentModelIds = new string[] { "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000001" };
            var modelId = Recording.GenerateId();

            RequestFailedException ex = Assert.ThrowsAsync<RequestFailedException>(async () => await client.ComposeDocumentModelAsync(WaitUntil.Started, fakeComponentModelIds, modelId));
            Assert.AreEqual("InvalidRequest", ex.ErrorCode);
        }

        #endregion

        #region Get

        [RecordedTest]
        [TestCase(true)]
        [TestCase(false)]
        public async Task GetDocumentModel(bool useTokenCredential)
        {
            var client = CreateDocumentModelAdministrationClient(useTokenCredential);
            var options = new BuildDocumentModelOptions()
            {
                Description = $"This model was generated by a .NET test.",
                Tags = { { "tag1", "value1" }, { "tag2", "value2" } }
            };
            await using var disposableModel = await BuildDisposableDocumentModelAsync(options: options);

            DocumentModelDetails expected = disposableModel.Value;
            DocumentModelDetails model = await client.GetDocumentModelAsync(disposableModel.ModelId);

            DocumentAssert.AreEqual(expected, model);
        }

        [RecordedTest]
        public void GetDocumentModelThrowsWhenModelDoesNotExist()
        {
            var client = CreateDocumentModelAdministrationClient();
            var fakeId = "00000000-0000-0000-0000-000000000000";

            RequestFailedException ex = Assert.ThrowsAsync<RequestFailedException>(async () => await client.GetDocumentModelAsync(fakeId));
            Assert.AreEqual("NotFound", ex.ErrorCode);
        }

        #endregion

        #region List

        [RecordedTest]
        [TestCase(true)]
        [TestCase(false)]
        public async Task GetDocumentModels(bool useTokenCredential)
        {
            var client = CreateDocumentModelAdministrationClient(useTokenCredential);
            var options = new BuildDocumentModelOptions()
            {
                Description = $"This model was generated by a .NET test.",
                Tags = { { "tag1", "value1" }, { "tag2", "value2" } }
            };

            // Make the models slightly different to make sure the cache won't return copies of the same model.
            await using var disposableModel0 = await BuildDisposableDocumentModelAsync(ContainerType.Singleforms, options);
            await using var disposableModel1 = await BuildDisposableDocumentModelAsync(ContainerType.MultipageFiles, options);

            var idMapping = new Dictionary<string, DocumentModelSummary>();
            var expectedIdMapping = new Dictionary<string, DocumentModelDetails>()
            {
                { disposableModel0.ModelId, disposableModel0.Value },
                { disposableModel1.ModelId, disposableModel1.Value }
            };

            await foreach (DocumentModelSummary model in client.GetDocumentModelsAsync())
            {
                if (expectedIdMapping.ContainsKey(model.ModelId))
                {
                    idMapping.Add(model.ModelId, model);
                }

                if (idMapping.Count == expectedIdMapping.Count)
                {
                    break;
                }
            }

            foreach (string id in expectedIdMapping.Keys)
            {
                Assert.True(idMapping.ContainsKey(id));

                DocumentModelSummary model = idMapping[id];
                DocumentModelDetails expected = expectedIdMapping[id];

                Assert.AreEqual(expected.ModelId, model.ModelId);
                Assert.AreEqual(expected.Description, model.Description);
                Assert.AreEqual(expected.ServiceVersion, model.ServiceVersion);
                Assert.AreEqual(expected.CreatedOn, model.CreatedOn);
                Assert.AreEqual(expected.ExpiresOn, model.ExpiresOn);

                CollectionAssert.AreEquivalent(expected.Tags, model.Tags);
            }
        }

        #endregion List

        #region Delete

        [RecordedTest]
        [TestCase(true)]
        [TestCase(false)]
        public async Task DeleteDocumentModel(bool useTokenCredential)
        {
            // Make sure we don't cache the model we'll delete, otherwise this may affect other tests.
            var client = CreateDocumentModelAdministrationClient(useTokenCredential);
            var disposableModel = await BuildDisposableDocumentModelAsync(skipCaching: true);

            var response = await client.DeleteDocumentModelAsync(disposableModel.ModelId);

            Assert.AreEqual((int)HttpStatusCode.NoContent, response.Status);
        }

        [RecordedTest]
        public void DeleteDocumentModelThrowsWhenModelDoesNotExist()
        {
            var client = CreateDocumentModelAdministrationClient();
            var fakeModelId = "00000000-0000-0000-0000-000000000000";

            RequestFailedException ex = Assert.ThrowsAsync<RequestFailedException>(async () => await client.DeleteDocumentModelAsync(fakeModelId));
            Assert.AreEqual("NotFound", ex.ErrorCode);
        }

        #endregion

        [RecordedTest]
        public async Task GetCopyAuthorization()
        {
            var client = CreateDocumentModelAdministrationClient();
            var modelId = Recording.GenerateId();

            DocumentModelCopyAuthorization copyAuthorization = await client.GetCopyAuthorizationAsync(modelId);

            Assert.AreEqual(modelId, copyAuthorization.TargetModelId);
            Assert.AreEqual(TestEnvironment.ResourceId, copyAuthorization.TargetResourceId);
            Assert.AreEqual(TestEnvironment.ResourceRegion, copyAuthorization.TargetResourceRegion);
            Assert.IsNotNull(copyAuthorization.AccessToken);
            Assert.IsNotEmpty(copyAuthorization.AccessToken);
            Assert.Greater(copyAuthorization.ExpiresOn, Recording.UtcNow);
        }
    }
}
